///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/ham/er_local.cpp
///
/// \brief   implementation header file for Lowdin (i.e. S^(-1/2)) orbital localization
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<string>
#include<vector>
#include<complex>

#include<formic/ham/lowdin_local.h>
#include<formic/ham/ham.h>
#include<formic/ham/integral_transform.h>
#include<formic/input/base.h>
#include<formic/exception.h>
#include<formic/output/output.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/lapack/interface.h>
#include<formic/all_data.h>
#include<formic/archive.h>
#include<formic/numeric/matrix.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   performs a Lowdin (i.e. S^(-1/2)) localization on the supplied orbitals
///
/// \param[in]     orb_inds     the indices of the orbitals to localize
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::Hamiltonian<S>::lowdin_localize(const std::vector<int> & orb_inds) {

  // get number of orbitals we are localizing
  const int nloc = orb_inds.size();

  // print some info
  if ( formic::mpi::rank() == 0 ) {
    formic::of << boost::format("   ==================================================") << std::endl;
    formic::of << boost::format("   ======     Lowdin S^(-1/2) Localization     ======") << std::endl;
    formic::of << boost::format("   ==================================================") << std::endl;
    formic::of << std::endl;
    formic::of << boost::format("localizing orbitals ");
    for (int i = 0; i < nloc; i++)
      formic::of << boost::format(" %2i") % orb_inds.at(i);
    formic::of << std::endl;
    formic::of << std::endl;
  }

  // check index sanity
  if ( ! formic::all_different_by_loops(orb_inds.begin(), orb_inds.end()) )
    throw formic::Exception("repeated orbital index in formic::Hamiltonian::lowdin_localize");
  if ( orb_inds.size() > _norb )
    throw formic::Exception("too many orbital indices in formic::Hamiltonian::lowdin_localize");
  for (int i = 0; i < nloc; i++) {
    if ( orb_inds.at(i) < 0 )
      throw formic::Exception("negative orbital index in formic::Hamiltonian::lowdin_localize");
    if ( orb_inds.at(i) >= _norb )
      throw formic::Exception("orbital index of %i is out of range in formic::Hamiltonian::lowdin_localize") % orb_inds.at(i);
  }

  // prepare a matrix to hold the localizing orbital transformation
  std::vector<S> tot_rot(_norb*_norb, formic::zero(S()));

  // compute the localizing transformation on root process
  if ( formic::mpi::rank() == 0 ) {

    // compute S^(-1/2) for the specified set of orbitals
    std::vector<S> s_neg_half(nloc*nloc, formic::zero(S()));
    for (int j = 0; j < nloc; j++)
    for (int i = 0; i < nloc; i++)
      s_neg_half[i+j*nloc] = _ovi[ orb_inds[i] + orb_inds[j] * _norb ];
    formic::matrix_neg_half(nloc, &s_neg_half.at(0));

    // compute the orbital rotation needed to orthonormalize the selected orbitals via S^(-1/2)
    std::vector<S> rot1(_norb*_norb, formic::zero(S()));
    for (int i = 0; i < _norb; i++)
      rot1.at(i+i*_norb) = formic::unity(S());
    for (int j = 0; j < nloc; j++)
    for (int i = 0; i < nloc; i++)
      rot1[ orb_inds[i] + orb_inds[j] * _norb ] = formic::conj( s_neg_half[i+j*nloc] );

    // compute a transformed overlap matrix
    std::vector<S> trn_ovi(_norb*_norb, formic::zero(S()));
    formic::transform_oei(_norb, &rot1.at(0), &_ovi.at(0), &trn_ovi.at(0));

    // print the intermediate overlap integrals
    this->print_oei(_norb, trn_ovi, "half-transformed overlap integrals:");

    // get the indices of the orbitals we are not localizing
    const std::vector<int> others = formic::other_indices(_norb, orb_inds);

    // compute the additional orbital rotation necessary to project the selected orbitals out of all other orbitals
    std::vector<S> rot2(_norb*_norb, formic::zero(S()));
    for (int i = 0; i < _norb; i++)
      rot2.at(i+i*_norb) = formic::unity(S());
    for (int j = 0; j < others.size(); j++)
    for (int i = 0; i < nloc; i++)
      rot2[ others[j] + orb_inds[i] * _norb ] -= trn_ovi[ orb_inds[i] + others[j] * _norb ];

    // compute the total orbital rotation
    formic::xgemm('N', 'N', _norb, _norb, _norb, formic::unity(S()), &rot2.at(0), _norb, &rot1.at(0), _norb, formic::zero(S()), &tot_rot.at(0), _norb);

  }

  // apply the total rotation to the hamiltonian
  this->transform_orbitals(&tot_rot.at(0));

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   performs a Lowdin (i.e. S^(-1/2)) localization on all the orbitals
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::Hamiltonian<S>::lowdin_localize() {

  // get a list of all orbitals
  std::vector<int> orb_inds(_norb, 0);
  for (int i = 0; i < _norb; i++)
    orb_inds.at(i) = i;

  // perform the localization
  this->lowdin_localize(orb_inds);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Lowdin_Localize_Select" command, which performs a Lowdin
///          orbital localization on the given subset of the orbitals.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::Lowdin_Localize_Select_command_function(formic::Archive & arch, formic::AllData & data) {

  // read list of orbitals to localize
  std::vector<int> orb_inds;
  formic::mpi::read_and_bcast(arch, orb_inds, "Failed to read indices of orbitals to localize in formic::Lowdin_Localize_Select_command_function.");

  // execute the localization
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    data.hd.lowdin_localize(orb_inds);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    data.hc.lowdin_localize(orb_inds);
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"Lowdin_Localize_Select\"") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Lowdin_Localize_All" command, which performs a Lowdin
///          orbital localization on all of the orbitals.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::Lowdin_Localize_All_command_function(formic::Archive & arch, formic::AllData & data) {

  // execute the localization
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    data.hd.lowdin_localize();
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    data.hc.lowdin_localize();
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"Lowdin_Localize_All\"") % data.userinp.get<std::string>("scalar_type");

}
